What does this PR do?

verl/models/mcore/model_forward_fused.py currently replaces the full GPTModel.forward through _fused_GPTModel_forward so the training loop can pass temperature into Megatron's _postprocess boundary, where linear_cross_entropy computes fused log-probabilities and entropy.

Megatron-LM PR #4686 added the native output_processor / output_processor_context extension point at that boundary. This PR uses that hook for verl's fused Megatron log-probability and entropy path while retaining the legacy monkey-patch as a compatibility fallback.

The fused forward mode is decided once at model build time:

• auto is the default. It uses the native hook when the actual bound Megatron model forward supports both output_processor and output_processor_context, and otherwise falls back to the legacy patch.
• hook requires the native hook and fails fast during model construction if the required contract is unavailable.
• legacy forces the existing monkey-patched GPTModel.forward path.

The selected mode is stored on the real Megatron model instance. Runtime forward calls read that fixed mode and do not repeatedly inspect the environment.

The live Megatron engine path is covered through fused_forward_model_engine. The legacy fused_forward_model caller is kept compatible as well. Megatron's build_schedule_plan() / 1F1B postprocess integration is outside this PR and remains part of the broader upstream work tracked in NVIDIA/Megatron-LM#4590.

Checklist Before Starting

• Search for similar PRs. Query: https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+megatron+output_processor
• The PR title follows the required [module] type: description format.
• I found no duplicate open PR implementing the same verl-side migration.
• I reviewed the relevant upstream and local context:
  • NVIDIA/Megatron-LM#4590
  • NVIDIA/Megatron-LM#4686
  • verl-project/verl#6927
  • verl-project/verl#6951

Test

Focused tests added in this revision

pytest -s -x tests/models/test_model_forward_fused.py

Result:

6 passed

The focused tests cover:

• auto selecting the hook path when the bound model forward supports the full hook contract.
• auto falling back to the legacy monkey-patch when the hook contract is unavailable.
• forced hook mode failing fast during model construction when required parameters are missing.
• invalid environment values raising ValueError.
• runtime callers using the mode fixed during model construction rather than rereading the environment.
• hook mode passing temperature through output_processor_context without forwarding temperature= to native Megatron forward.
• legacy mode continuing to pass temperature= to the patched forward.
• sequence-parallel gather occurring before the fused linear cross-entropy call.
• tied and untied output-weight resolution.
• patch and unpatch behavior remaining idempotent.

Pre-submit checks:

pre-commit run --all-files

Result:

passed

git diff upstream/main...HEAD --check

Result:

clean

Previously completed validation

The numerical and live integration validation below was completed before the latest rebase and was not rerun in full for this revision.

Environment used for the same-batch numerical validation:

Megatron-Core: 0.18.0+97f3bce88
PyTorch: 2.9.1+cu128
Transformer Engine: 2.10
flash-attn: 2.8.3
dtype: bf16

A separate 4×A100 five-step GRPO smoke previously exercised:

• Megatron engine with TP=2 and PP=1.
• vLLM rollout.
• fused log-probability and entropy computation.
• parameter, gradient, and optimizer offload.
• optimizer updates.
• rollout-to-actor weight resharding.
• distributed checkpoint save/load.

Both the hook and legacy paths completed that five-step integration run. The two independent runs saw different sampled trajectories, so per-step metrics were not expected to be identical; exact numerical equivalence was established by the same-weight, same-batch A/B validation above.

Forward memory usage was also measured as identical between the hook and legacy paths (Δ = 0).

Rebased local smoke limitation

A rebased local smoke was attempted, but the checkout runtime did not match the project's full CUDA image dependencies and the run stopped before Megatron model construction. The output_processor path was therefore not reached in that local run.

The focused tests above pass, and the previously completed same-batch numerical validation and live five-step integration validation remain passing. The official project CI environment will provide the final rebased integration check.

Known limitation

The FP8 padding path was not exercised:

use_fp8_padding
config.fp8 in {e4m3, hybrid}

All numerical validation reported above used bf16. The callback itself does not introduce FP8-specific logic, but this PR does not claim coverage of the FP8 input-padding path.

API and Usage Example

# Recommended default: use the native hook when supported,
# otherwise fall back to the legacy forward patch.
VERL_FUSED_USE_OP_HOOK=auto

# Require the native hook and fail fast if unsupported.
VERL_FUSED_USE_OP_HOOK=hook

# Force the legacy GPTModel.forward patch.
VERL_FUSED_USE_OP_HOOK=legacy

Leaving VERL_FUSED_USE_OP_HOOK unset or setting it to an empty value is equivalent to auto.

Accepted aliases are:

auto:   unset, "", auto
hook:   1, true, yes, hook
legacy: 0, false, no, legacy

Any other value raises ValueError during model construction.

Design & Code Changes

• Parse VERL_FUSED_USE_OP_HOOK once during patch_fused_forward.
• Resolve the actual underlying Megatron GPT model before capability probing.
• Probe the actual bound model.forward signature for both output_processor and output_processor_context.
• Store the selected mode on the real model instance.
• Have both fused forward callers read the stored mode instead of rereading the process environment.
• Keep hook and legacy selection stable even if the environment variable changes after model construction.
• In hook mode, leave native GPTModel.forward untouched and do not create forward_backup.
• In legacy mode, retain the existing monkey-patch and forward_backup compatibility path.
• Keep repeated patch and unpatch calls safe and idempotent.
• Raise a clear RuntimeError when forced hook mode is requested but the bound forward lacks the required hook contract.
• Remove the redundant output-weight assertion identified during review.

Two correctness details are intentional:

1. Megatron passes output_processor_context to the callback using the keyword context. The callback parameter is therefore named context, while the native forward call correctly uses output_processor_context=....

2. The hook's output-weight contract differs from the legacy copied forward:

   • tied embeddings provide the shared weight through output_weight;
   • untied embeddings provide output_weight=None, and the callback uses output_layer.weight.

The callback therefore resolves the fused-kernel weight with:

weight = output_weight if output_weight is not None else output_layer.weight

Under sequence parallelism, the callback gathers hidden_states before invoking linear_cross_entropy, matching the established legacy behavior.

Checklist Before Submitting

• I have read the contribution guide.
• pre-commit run --all-files passed.
• Focused tests were added.
• The CI workflow was updated to run the focused tests.
• Separate documentation was not added; the environment-variable API is documented in the usage example above.
• CI has not yet been requested in the project's ci-request channel.
• Not applicable: this PR does not modify the recipe submodule.

AI assistance disclosure: OpenAI Codex assisted with implementation and test execution; I reviewed the resulting changes and test outputs.